#include "../include/association.h"


using namespace hitcrt;

/**
 * @brief 设置代价矩阵
 * @param data_ptr,数据指针,类型为float,似乎不能使用常指针
 * @note 这里用到了数组赋值，参考https://blog.csdn.net/sinat_28751869/article/details/79425491
 */ 
void KMAssociation::set_cost_matrix(float *data_ptr,int n,int m)
{
    // this->cost_matrix.resize(n,m);
    // this->mask = Eigen::MatrixXi::
    this->cost_matrix = Eigen::Map<Eigen::Matrix<float,Eigen::Dynamic,Eigen::Dynamic,Eigen::RowMajor> >(data_ptr,n,m);
}

bool KMAssociation::dfs(int u,int cols)
{ //捕获所有外部变量和this指针
    this->visx[u] = 1;//标记进入增广路
    for(int v = 0; v < cols; v++)
    {
        if(!this->visy[v] && this->cost_matrix(u,v) != INF)//如果Y部的点还没进入增广路,并且存在路径
        // if(!this->visy[v])
        {
            float t = this->wx[u] + this->wy[v] - this->cost_matrix(u,v);
                if(fabs(t) < 1e-5 )//t为0说明是相等子图
                {
                    this->visy[v] = true;//加入增广路

                    //如果Y部的点还未进行匹配
                    //或者已经进行了匹配，可以从原来的匹配反向找到增广路
                    //那就可以进行匹配
                    if(this->cy[v] == -1 || dfs(cy[v],cols))
                    {
                        this->cx[u] = v;
                        this->cy[v] = u;//进行匹配
                        return true;
                    }
                }
                else if(t > 0)//此处t一定是大于0，因为顶标之和一定>=边权
                {
                    slack[v] = std::min(slack[v], t);
                    //slack[v]存的是Y部的点需要变成相等子图顶标值最小增加多少
                }
            }
        }
        return false;
    };

/**
 * @brief 数据关联 
 * @param row_indices,输出匹配成功的列索引，对应tracking
 * @param col_indices,输出匹配成功的行索引，对应detection，其顺序和row的对应
 * @param maximize,是否计算最大权重匹配,暂时没有用
 * @ref https://blog.csdn.net/lemonxiaoxiao/article/details/108704280
 * @ref https://zhuanlan.zhihu.com/p/89380238
 * @note 源代码是求完备匹配，必须要求为方阵，这里我修改了一下，直接求最大匹配
 * @note 其实将cost matrix改为方阵也是可以的
 */ 
int KMAssociation::associate(std::vector<int> &row_indices,std::vector<int> &col_indices,bool maximize)
{
#ifdef ASSOCIATION_DEBUG
    std::cout << "cost matrix = \n"<<this->cost_matrix<<std::endl;
#endif
    row_indices.clear();
    col_indices.clear();

    int cols = this->cost_matrix.cols();
    int rows = this->cost_matrix.rows();
    int maxn = (cols > rows) ? cols:rows;

    this->wx = new float[rows];//左部点权重
    this->wy = new float[cols];//右部点权重

    this->visx = new bool[rows];//访问标记,是否在交错树中
    this->visy = new bool[cols];

    this->cx = new int[rows];//与右侧点匹配的左点
    this->cy = new int[cols];//与左侧点匹配的右点

    this->slack = new float[maxn];//边权和顶标最小的差值

    //初始化权值
    memset(cx,-1,sizeof(int)*rows);
    memset(cy,-1,sizeof(int)*cols);
    memset(wx,0,sizeof(float)*rows);
    memset(wy,0,sizeof(float)*cols);
    for(int i = 0; i < rows; i++)//预处理出顶标值
    {
        for(int j = 0; j < cols; j++)
        {
            if(this->cost_matrix(i,j) >= INF)continue;
            wx[i] = std::max(wx[i], this->cost_matrix(i,j));
        }
    }

    for(int i = 0; i < rows; i++)//枚举X部的点
    {
        for (int index=0;index<maxn;index++)
        {
           slack[index] = INF;
        }
        while(1)
        {

            memset(visx, 0, sizeof(bool)*rows);
            memset(visy, 0, sizeof(bool)*cols);
            if(this->dfs(i,cols))break;//已经匹配正确

            float minz = INF;
            for(int j = 0; j < cols; j++)
                if(!visy[j] && minz > slack[j])
                    //找出还没经过的点中，需要变成相等子图的最小额外增加的顶标值
                    minz = slack[j];
            //和全局变量不同的是，全局变量在每次while循环中都需要赋值成INF，每次求出的是所有点的最小值
            //而slack数组在每个while外面就初始化好，每次while循环slack数组的每个值都在用到
            //在一次增广路中求出的slack值会更准确，循环次数比全局变量更少


            //还未匹配，将X部的顶标减去minz，Y部的顶标加上minz
            for(int j = 0; j < rows; j++)
                if(visx[j])wx[j] -= minz;
            for(int j = 0; j < cols; j++)
                //修改顶标后，要把所有不在交错树中的Y顶点的slack值都减去minz
                if(visy[j])wy[j] += minz;
                else slack[j] -= minz;
#ifdef ASSOCIATION_DEBUG
            std::cout<<"[debug message]:"<<i<<std::endl;
            std::cout<<"\t slack : ";
            for (int index=0;index<maxn;index++)
            {
                std::cout<<slack[index]<<"\t";
            }
            std::cout<<std::endl;
            std::cout<<"\t wx:";
            for (int index=0;index<rows;index++)
            {
                std::cout<<wx[index]<<"\t";
            }
            std::cout<<std::endl;
            std::cout<<"\t wy:";
            for (int index=0;index<cols;index++)
            {
                std::cout<<wy[index]<<"\t";
            }
            std::cout<<std::endl;
#endif
        }
    }
    float ans = 0;//二分图最优匹配权值
    for(int i = 0; i < rows; i++)
    {
        if(cx[i] != -1)
        {
            if (this->cost_matrix(i,cx[i]) >=INF){continue;}
            ans += this->cost_matrix(i,cx[i]);
            row_indices.push_back(i);//加入匹配的行索引
            col_indices.push_back(cx[i]);//加入匹配的列索引
        }
    }
        

    delete[] wx;//左部点权重
    delete[] wy;//右部点权重
    delete[] visx;//访问标记,是否在交错树中
    delete[] visy;
    delete[] cx;//与右侧点匹配的左点
    delete[] cy;//与左侧点匹配的右点
    delete[] slack;//边权和顶标最小的差值

    wx = nullptr;
    wy = nullptr;
    visx = nullptr;
    visy = nullptr;
    cx = nullptr;
    cy = nullptr;
    slack = nullptr;

    return ans;
}
